"""Tools for manipulating processes, mutexes, channels, and selects in CIRCUIT.

"""

# TODO(chris): flesh out description above.

from __future__ import division, print_function

import errno
import json
import multiprocessing
import os
import os.path
import re
import sys
import threading
import time
import Queue

class Error(Exception): pass
class RippleCreationFailed(Error): pass
class RippleExists(Error): pass
class NoSuchRipple(Error): pass
class InvalidChanStat(Error): pass
class ChannelCapacityNotSet(Error): pass

SPLIT_RE = re.compile(r":\s*")

class _Ripple(object):
  """Base class for ripples. Contains fundamental behavior."""

  def __init__(self, kind, slab, name=None):
    """Create a ripple wrapper, and possibly a new underlying ripple."""
    self.ripple_dir = RippleDirectory(slab, kind)
    if name is None:
      name = self.ripple_dir.create_name()

    self.slab = self.ripple_dir.slab
    self.kind = self.ripple_dir.kind
    self.name = name
    self.path = self.ripple_dir.find_name(name)
    if self.path is None:
      raise NoSuchRipple(self.ripple_dir.name_path(name))

  def __repr__(self):
    return "{kind}({slab!r}, {name!r})".format(kind=self.kind, slab=self.slab, name=self.name)

  def help(self):
    return open(os.path.join(self.path, "help")).read()

class Mutex(_Ripple):
  """Track and manipulate a circuit mutex.

  Example use:

  m = Mutex("/path/to/circuit/lang", "mymutex")

  def work():
    print("Got the lock - go to work")

  success = m.acquire_and_do(work, 5)  # 5-second timeout.
  """

  def __init__(self, slab, name=None):
    """Create a new mutex helper.

    Args:
      slab: a path to the slab containing the desired Mutex directory.
      name: the mutex name to assign (or use). If None, creates a new unique
          name and uses that.
    """
    super(Mutex, self).__init__("Mutex", slab, name)
    self._waiter = FileWriteWaiter(self.lock_path, "")

  def copy(self):
    return self.__class__(self.slab, name=self.name)

  def acquire_and_do(self, f, timeout=None):
    """Acquire the mutex and run function f before releasing it.

    Args:
      f: run this function while holding the lock
      timeout: stop waiting after timeout and return False. None = forever.

    Returns:
      True if the function is run, False otherwise.
    """
    result = self._waiter.do(f, timeout)
    return True if result is not None else False

  def abandon(self):
    """Abandon this mutex - stop waiting for acquisition.

    Returns:
      True if the mutex was abandoned, false if it is acquired and busy.
    """
    return self._waiter.abandon()

  def stat(self):
    """Return stats on this mutex (num holding, num waiting)."""
    vals = {}
    with open(self.stat_path, "rb") as f:
      for line in f:
        parts = SPLIT_RE.split(line.rstrip(), 1)
        vals[parts[0].lower()] = int(parts[1])
    return vals["numholding"], vals["numwaiting"]

  def reset(self):
    """Reset the mutex, cause all waiters to drop out."""
    with open(self.reset_path, "ab") as f:
      print("reset", file=f)

  @property
  def lock_path(self):
    return os.path.join(self.path, "lock")

  @property
  def reset_path(self):
    return os.path.join(self.path, "reset")

  @property
  def stat_path(self):
    return os.path.join(self.path, "stat")



class Proc(_Ripple):
  """Manage a process ripple.

  Example use:

  >>> p = Proc("/path/to/slab/lang")
  >>> p.set_spec("/bin/bash", arg=["-c", 'echo "Hello, ${OBJ}!"'], env=["OBJ=World"])
  >>> p.start()
  >>> p.wait(timeout=5)
  >>> print(open(p.stdout_path).read())
  """

  def __init__(self, slab, name=None):
    """Create a new process helper.

    Args:
      slab: a path to the slab containing the desired Proc directory.
      name: the process name to assign (or use). If None, creates a new unique
          name and uses that.
    """
    super(Proc, self).__init__("Proc", slab, name)
    self._waiter = FileReadWaiter(os.path.join(self.path, "wait"))

  def copy(self):
    return self.__class__(self.slab, name=self.name)

  def spec(self):
    """Return the spec as the binary, arguments, and env list.

    Returns:
      Command string, list of arg strings, and list of env strings.
    """
    spec = json.load(open(os.path.join(self.path, "spec_json")))
    return spec.get("bin", []), spec.get("arg", []), spec.get("env", [])

  def set_spec(self, cmd, arg=(), env=()):
    """Set the process spec if possible.

    Args:
      cmd: a string indicating which binary to run
      arg: a list of argument strings
      env: a list of "VAR=something" strings
    """
    json.dump({"bin": cmd,
               "arg": arg,
               "env": env},
              open(os.path.join(self.path, "spec_json"), "ab"))

  def stat(self):
    """Return the process ID and status, if available.

    Returns:
      PID (int or None), status (str or None)
    """
    pid = None
    status = None
    with open(os.path.join(self.path, "stat")) as f:
      for line in f:
        parts = line.split(":", 1)
        key = parts[0].lower().strip()
        if key == "pid":
          pid = parts[1].strip()
          if pid:
            pid = int(val)
          else:
            pid = None
        elif key == "status":
          status = parts[1].strip()
    return pid, status

  def wait(self, timeout=None):
    """Wait for the process to complete, and return the value (or None if timed out)."""
    return self._waiter.do(timeout=timeout)

  def abandon_wait(self):
    """Cause the wait() call to finish early."""
    return self._waiter.abandon()

  def _ctrl(self, *args):
    with open(os.path.join(self.path, "ctrl"), "ab") as f:
      print(" ".join(args), file=f)

  def start(self):
    """Start the process."""
    self._ctrl("start")

  def kill(self, signal='9'):
    """Kill the process with the given signal.

    Args:
      signal: when converted into a string, can be of the form '-9' or 'HUP' or
          '-INT', etc.
    """
    self._ctrl("kill", str(signal))

  def reset(self):
    """Reset the process, including all arguments, etc."""
    self._ctrl("reset")

  @property
  def stderr_path(self):
    return os.path.join(self.path, "stderr")

  @property
  def stdout_path(self):
    return os.path.join(self.path, "stdout")

  @property
  def stdin_path(self):
    return os.path.join(self.path, "stdin")

class Chan(_Ripple):
  """Manage a circuit channel.

  Example Use:

  >>> c = Chan("/path/to/slab/lang", "mychan")
  >>> c.set_cap(1)
  >>> c.send("hi there!\\n")
  >>> print(c.recv())

  Of course, you don't get the full benefit unless you send and receive from
  different processes. Simply create a channel with the same name in another
  interpreter and they will communicate back and forth.
  """

  class Stat(object):
    """Status information for a channel."""

    def __init__(self, capacity=None, in_bytes=0, in_count=0,
                 out_bytes=0, out_count=0, closed=False):
      self.capacity = capacity
      self.in_bytes = in_bytes
      self.out_bytes = out_bytes
      self.in_count = in_count
      self.out_count = out_count
      self.closed = closed

    def __repr__(self):
      fields = ["capacity",
                "in_bytes",
                "in_count",
                "out_bytes",
                "out_count",
                "closed"]
      pieces = ["{}={!r}".format(x, getattr(self, x)) for x in fields]
      return "{}({})".format(self.__class__.__name__, ", ".join(pieces))

  def __init__(self, slab, name=None):
    """Create a new channel helper.

    Args:
    slab: a path to the slab containing the desired Proc directory.
    name: the channel name to assign (or use). If None, creates a new unique
      name and uses that.
    """
    super(Chan, self).__init__("Chan", slab, name)

    self._recv_waiter = FileReadWaiter(self.recv_path)
    self._send_waiter = FileWriteWaiter(self.send_path)

  def copy(self):
    return self.__class__(self.slab, name=self.name)

  def stat(self):
    """Return a Chan.Stat object with status info for this channel."""
    headers = {
      "capacity": ("capacity", lambda s: None if s.lower().strip() == "none" else int(s)),
      "byteswritten": ("in_bytes", int),
      "bytesread": ("out_bytes", int),
      "recordswritten": ("in_count", int),
      "recordsread": ("out_count", int),
      "closed": ("closed", lambda s: {"false": False, "true": True}.get(s))
    }

    args = {}

    with open(os.path.join(self.path, "stat")) as f:
      for line in f:
        parts = SPLIT_RE.split(line.rstrip(), 1)
        varname, converter = headers.get(parts[0].lower(), (None, None))
        if varname is None:
          raise InvalidChanStat("Unknown field in channel stat line: {!r}".format(line))
        args[varname] = converter(parts[1].lower())

    return self.Stat(**args)

  def close(self):
    with open(os.path.join(self.path, "close"), "ab") as f:
      print("close", file=f)

  def set_cap(self, val):
    """Set the channel's capacity to a non-negative value."""
    with open(os.path.join(self.path, "cap"), "ab") as f:
      print(max(0, val), file=f)

  def cap(self):
    with open(os.path.join(self.path, "cap")) as f:
      return int(f.read())

  def tryrecv(self):
    """Try to receive on this channel, non-blocking.

    Returns:
      None if not ready, a (possibly empty) byte string if ready.
    """
    if self.cap() < 0:
      raise ChannelCapacityNotSet()

    try:
      with open(self.tryrecv_path, "rb") as f:
        return f.read()
    except IOError, e:
      if e.errno != errno.EBUSY:
        raise e
      return None

  def trysend(self, val):
    """Try to send the given byte string, non-blocking.

    Returns:
      True if sent, False if not ready.
    """
    if self.cap() < 0:
      raise ChannelCapacityNotSet()

    try:
      with open(self.trysend_path, "ab") as f:
        f.write(val)
        return True
    except IOError, e:
      if e.errno != errno.EBUSY:
        raise e
      return False

  def recv(self, timeout=None):
    """Receive on this channel, block with an optional timeout.

    Args:
      timeout: time to wait, or None if infinite.

    Returns:
      None if wait failed, else a possibly empty bytestring.
    """
    if self.cap() < 0:
      raise ChannelCapacityNotSet()
    return self._recv_waiter.do(timeout=timeout)

  def abandon_recv(self):
    """Abandon any in-process receives."""
    return self._recv_waiter.abandon()

  def send(self, data=b"", timeout=None):
    """Send data on this channel, block with optional timeout.

    Args:
      data: a (possibly empty) bytestring to send.
      timeout: time to wait (sec) or None if infinite.

    Returns:
      False if timed out, True if sent.
    """
    if self.cap() < 0:
      raise ChannelCapacityNotSet()
    return self._send_waiter.do(data=data, timeout=timeout)

  def abandon_send(self):
    "Abandon any in-process sends."""
    return self._send_waiter.abandon()

  @property
  def send_path(self):
    return os.path.join(self.path, "send")

  @property
  def recv_path(self):
    return os.path.join(self.path, "recv")

  @property
  def trysend_path(self):
    return os.path.join(self.path, "trysend")

  @property
  def tryrecv_path(self):
    return os.path.join(self.path, "tryrecv")

class Select(_Ripple):
  """Manage a select ripple.

  Example use (assumes channels called "c1" and "c2"):

  >>> s = ripple.Select("/path/to/slab/lang")
  >>> s.set_cases([(c1.send_path, "hi there!\\n"),
                   (c2.recv_path, None), ripple.Select.DefaultCase])
  >>> s.select(5)
  """

  DefaultCase = ("default", "")

  def __init__(self, slab, name=None):
    super(Select, self).__init__("Select", slab, name)
    self._waiter = FileReadWaiter(os.path.join(self.path, "select"))

  def copy(self):
    return self.__class__(self.slab, self.name)

  def set_cases(self, cases):
    """Set the cases for this select statement.

    Args:
      cases: a list of (channel_file_name, data) pairs. The data can be None
          for receive channels and a byte string for send channels. The case
          can also be Select.DefaultCase.
    """
    case_obj = [{"chan": cname, "data": data} for cname, data in cases]
    json.dump(case_obj, open(os.path.join(self.path, "cases_json"), "ab"))

  def cases(self):
    """Get the cases for this select as a list of (name, data) pairs."""
    case_obj = json.load(open(os.path.join(self.path, "cases_json"), "rb"))
    return [(c["chan"], c["data"]) for c in case_obj]

  def select(self, timeout=None):
    """Perform a select, blocking with optional timeout.

    To make this non-blocking, simply include a default case in set_cases.

    Args:
      timeout: optional seconds to time out. None is infinite.

    Returns:
      (chosen channel index, chosen name, [data or None]) or None
    """
    val = self._waiter.do(timeout=timeout)
    if val is None:
      return None

    parts = val.split(":", 1)
    idx = int(parts[0])
    parts = parts[1].split("\n", 1)
    return idx, parts[0], parts[1]

  def abandon_select(self):
    return self._waiter.abandon()

  def interrupt(self):
    with open(os.path.join(self.path, "interrupt"), "ab") as f:
      print("interrupt", file=f)

class RippleDirectory(object):
  def __init__(self, slab, kind):
    """Manage a ripple directory, e.g., something/lang/Mutex.

    Args:
      path: path to the slab (e.g., stuff/lang/myslab)
      kind: kind (str) of ripple, i.e., mutex, proc, select, or chan.
    """
    self.slab = slab
    self.kind = kind.capitalize()
    self.path = os.path.join(self.slab, self.kind)

  def slab_exists(self):
    return os.path.isdir(self.slab)

  def name_exists(self, name):
    return os.path.isdir(self.name_path(name))

  def name_path(self, name):
    return os.path.join(self.path, name)

  def create_name(self, name=None):
    """Create a new ripple. If no name, generate a unique one and return it.

    Args:
      name: Try to create a ripple of this name. If None, generate one.

    Raises:
      RippleExists: if trying to create a name that is already there.
      RippleCreationFailed: if something else goes wrong with creation.
    """
    if name is None:
      name = open(os.path.join(self.path, "new")).read().rstrip()
      if not self.name_exists(name):
        raise RippleCreationFailed("Could not create path {!r}".format(self.name_path(name)))
    else:
      if self.name_exists(name):
        raise RippleExists("Ripple already exists at {!r}".format(self.name_path(name)))
      try:
        os.mkdir(self.name_path(name))
      except OSError, e:
        raise RippleCreationFailed("Could not create path {!r}: {}".format(self.name_path(name, e)))
    return name

  def find_name(self, name):
    """Find the full path for the given name, if it exists, else None."""
    if not self.name_exists(name):
      return None
    return self.name_path(name)

  def find_or_create_name(self, name):
    """Find the given name or try to create it. Return the full path."""
    if not self.name_exists(name):
      self.create_name(name)
    return self.name_path(name)

  def help(self):
    return open(os.path.join(self.path, "help")).read()

  def __str__(self):
    return self.path


class FileOpWaiter(object):
  def __init__(self, path):
    self.path = path

    self._do_lock = threading.Lock()
    self._run_lock = threading.Lock()
    self._var_lock = threading.Lock()

    self._reset()

  def do_wait(self, ready_q, done_q, unused_data):
    try:
      with open(self.path, "rb") as f:
        ready_q.put((True, f.read()))
        done_q.get()
    except IOError, e:
      ready_q.put((False, e))
      done_q.get()

  def _reset(self):
    with self._var_lock:
      self._proc = None
      self._ready_q = None
      self._done_q = None

  def do(self, f=None, timeout=None, data=b""):
    """Run f after successfully opening the blocking resource. Optionally send data.

    Args:
      f: the function to run when the resource is acquired
      timeout: optional timeout (sec), None means infinity
      data: optional data to send for write operations

    Returns:
      None if the operation is interrupted, the data read or written if completed.
    """
    with self._do_lock:
      val = None

      with self._var_lock:
        self._ready_q = multiprocessing.Queue(1)
        self._done_q = multiprocessing.Queue(1)
        self._proc = multiprocessing.Process(target=self.do_wait,
                                             args=(self._ready_q, self._done_q, data))

      self._proc.start()
      try:
        success, val = self._ready_q.get(True, timeout)
        if not success:
          self._done_q.put(False)
          raise val
      except (Queue.Empty, IOError), e: # also raised by get if it times out
        if isinstance(e, IOError) and e.errno not in [errno.EBADF, errno.EBUSY]:
          print("Unexpected exception in waiting process:\n{}".format(e), file=sys.stderr)
        with self._var_lock:
          self._proc.terminate()
          self._proc.join()
          self._proc = None
        return None

      with self._run_lock:
        if f is not None:
          f()

      self._done_q.put(True) # non-blocking - queue has length 1

      with self._var_lock:
        self._proc.join()
        self._proc = None

      return val

  def abandon(self):
    with self._var_lock:
      if self._proc is None:
        return False

    with self._run_lock:
      with self._var_lock:
        if self._ready_q is not None:
          self._ready_q.put_nowait((False, None))
          return True

    return False

class FileReadWaiter(FileOpWaiter): pass

class FileWriteWaiter(FileOpWaiter):
  def __init__(self, path, data=b""):
    super(FileWriteWaiter, self).__init__(path)
    self.data = data

  def do_wait(self, ready_q, done_q, data=b""):
    try:
      with open(self.path, "ab") as f:
        f.write(data)
        ready_q.put((True, data))
        done_q.get()
    except IOError, e:
      ready_q.put((False, e))
      done_q.get()
